## helper functions to run other tests

import unittest, sys, pdb, traceback, getopt, os, math

from pyblog.misc.util import POSINF

class MyTestCase(unittest.TestCase):
  def assertPctEqual(self, actual, expected, pct, name):
    self.assertTrue(abs(float(actual)/expected - 1.0) <= float(pct)/100.0,
                    "%s not within %f pct actual=%f expected=%f %%err=%f" \
                    % (name, pct, actual, expected,
                       abs(float(actual)/expected - 1.0) * 100))
  
  def assertAbsEqual(self, actual, expected, absdiff, name):
    self.assertTrue(abs(actual - expected) <= abs(absdiff),
                    "%s not within %f : actual=%f expected=%f" \
                    % (name, abs(absdiff), actual, expected))
    
def parse_options():
  """
  parse the command-line arguments and return as a tuple
  (verbose, stats, trace, debug, args)
  """
  USAGE = """\
Usage: %s [options] [test] [...]

Options:
  -h, --help       Show this message
  -v, --verbose    Verbose test output
  -t, --trace      Enable PyBLOG tracing
  -d, --debug      Debug mode testing

Examples:
  %s                      - run all test cases/ test suites
  %s <test1> <test2>...   - run specific test cases/test suites
"""
  # default options
  ver = 0
  dispstats = False
  disptrace = False
  debug = False
  
  # parse command line
  options, args = getopt.getopt(sys.argv[1:], 'hvsdt',
                                ['help', 'verbose', 'stats', 'debug', 'trace'])
  
  # command line switches
  for opt, value in options:
    if opt in ('-h', '--help'):
      print USAGE % (sys.argv[0], sys.argv[0], sys.argv[0])
      sys.exit(0)
      
    if opt in ('-v', '--verbose'):
      ver = 2
      dispstats = True

    if opt in ('-d', '--debug'):
      debug = True

    if opt in ('-t', '--trace'):
      disptrace = True

  return [ver, dispstats, disptrace, debug, args]

def testmodule(mod):
  """
  Tests all the testcases in the module as per the command line arguments
  
  mod - module to test
  """
  _testmodule(mod, *parse_options())

def _testmodule(mod, ver, dispstats, disptrace, debug, args):
  mod.dispstats, mod.disptrace = dispstats, disptrace
  # now load the test names
  if not args:
    suite = unittest.TestLoader().loadTestsFromModule(mod)
  else:
    suite = unittest.TestLoader().loadTestsFromNames(args, mod)

  if not suite.countTestCases():
    return
  
  runner = unittest.TextTestRunner(verbosity=ver)
  
  if debug:
    try:
      suite.debug()
    except:
      traceback.print_exc()
      pdb.post_mortem(sys.exc_traceback)
  else:
    runner.run(suite)

def testdir(dirname):
  # get the user options and extract the module names (if any)
  opts = parse_options()
  modnames = opts[-1]
  opts[-1] = []
  if modnames:
    for modname in modnames:
      try:
        mod = __import__(modname)
      except:
        print "Skipping:", modname
        continue
      print "Testing:", modname
      _testmodule(mod, *opts)
  else:
    for filename in os.listdir(dirname):
      if filename[-3:].lower() == '.py':
        try:
          mod = __import__(filename[:-3])
        except:
          print "Skipping:", filename
          continue
        print "Testing:", filename
        _testmodule(__import__(filename[:-3]), *opts)

def gauss_density(theta, pt):
  """
  Measures the gaussian density at a point
  theta - pair of mean, sigma
  pt    - coordinate of point at which to compute density
  """
  return math.sqrt(1.0 / (2.0 * math.pi)) * (1.0 / float(theta[1])) \
                * math.exp(- 0.5 * (theta[0] - pt) ** 2 / float(theta[1]) ** 2)
  
def gauss_kl_diver(theta1, theta2):
  """
  theta - array of mean, sigma pairs
  theta1 - true parameters
  theta2 - guessed parameters

  We are assuming that each of the gaussians have the same weight
  """
  # We need to compute xmin and xmax the minimum and the maximum points to
  # measure divergence at. For these we'll use 10 times the standard deviation
  # Also, we need the step size for which we'll use half the standard dev.
  xmin = min(mean - 10.0 * sigma for mean,sigma in theta1)
  xmax = max(mean + 10.0 * sigma for mean,sigma in theta1)
  step = min(sigma/2.0 for _,sigma in theta1 + theta2)
  cnt = int((xmax - xmin)/step) + 1     # number of steps to take

  div = 0.0
  
  for i in range(cnt):
    x = xmin + step*i
    p1 = sum(gauss_density(t,x) for t in theta1) / float(len(theta1))
    p2 = sum(gauss_density(t,x) for t in theta2) / float(len(theta2))

    if p1==0:
      continue
    
    if p2 == 0.0:
      div += step * p1 * (math.log(p1) + 300)
    else:
      div += step * p1 * (math.log(p1) - math.log(p2))

  assert(div >= 0)

  return div

def gauss_kl_diver_wt(theta1, theta2):
  """
  theta - array of ((mean, sigma), wt)
  theta1 - true parameters
  theta2 - guessed parameters

  We are assuming that the weights add upto 1
  """
  # We need to compute xmin and xmax the minimum and the maximum points to
  # measure divergence at. For these we'll use 10 times the standard deviation
  # Also, we need the step size for which we'll use half the standard dev.
  xmin = min(mean - 10.0 * sigma for (mean,sigma),wt in theta1)
  xmax = max(mean + 10.0 * sigma for (mean,sigma),wt in theta1)
  step = min(sigma/2.0 for (mean,sigma),wt in theta1 + theta2)
  cnt = int((xmax - xmin)/step) + 1     # number of steps to take

  # verify that theta1 and theta2 are both normalized
  assert(abs(sum(wt for t,wt in theta1) - 1) < .001)
  assert(abs(sum(wt for t,wt in theta2) - 1) < .001)
    
  div = 0.0
  
  for i in range(cnt):
    x = xmin + step*i
    p1 = sum(gauss_density(t,x) * wt for t,wt in theta1)
    p2 = sum(gauss_density(t,x) * wt for t,wt in theta2)

    if p1==0:
      continue
    
    if p2 == 0.0:
      div += step * p1 * (math.log(p1) + 300)
    else:
      div += step * p1 * (math.log(p1) - math.log(p2))

  assert(div >= 0)

  return div
    
if __name__ == "__main__":
  # test the modules in the current directory
  testdir(os.path.dirname(os.path.abspath(sys.argv[0])))
  
